import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import tensorflow as tf
from tensorflow import keras
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score
import os

# 设置随机种子以确保结果可重现
np.random.seed(42)
tf.random.set_seed(42)

# 创建输出目录
os.makedirs('output', exist_ok=True)

# 加载数据
print("加载数据...")
data = pd.read_csv('data/completion_data.csv')

# 显示数据基本信息
print("\n数据基本信息:")
print(data.info())
print("\n数据前5行:")
print(data.head())
print("\n数据统计摘要:")
print(data.describe())

# 检查缺失值
print("\n检查缺失值:")
print(data.isnull().sum())

# 探索性数据分析
print("\n进行探索性数据分析...")

# 1. 分析星期几与完成率的关系
plt.figure(figsize=(10, 6))
sns.boxplot(x='day_of_week', y='completion_rate', data=data)
plt.title('星期几与完成率的关系')
plt.xlabel('星期几 (1=周一, 7=周日)')
plt.ylabel('完成率')
plt.savefig('output/day_of_week_vs_completion_rate.png')

# 2. 分析是否假日与完成率的关系
plt.figure(figsize=(10, 6))
sns.boxplot(x='is_holiday', y='completion_rate', data=data)
plt.title('是否假日与完成率的关系')
plt.xlabel('是否假日 (0=否, 1=是)')
plt.ylabel('完成率')
plt.savefig('output/is_holiday_vs_completion_rate.png')

# 3. 分析计划任务数量与完成率的关系
plt.figure(figsize=(10, 6))
sns.scatterplot(x='planned_tasks', y='completion_rate', data=data)
plt.title('计划任务数量与完成率的关系')
plt.xlabel('计划任务数量')
plt.ylabel('完成率')
plt.savefig('output/planned_tasks_vs_completion_rate.png')

# 4. 相关性分析
plt.figure(figsize=(10, 8))
correlation_matrix = data.drop('date', axis=1).corr()
sns.heatmap(correlation_matrix, annot=True, cmap='coolwarm', fmt='.2f')
plt.title('特征相关性矩阵')
plt.savefig('output/correlation_matrix.png')

# 准备数据用于模型训练
print("\n准备数据用于模型训练...")

# 选择特征和目标变量
X = data[['day_of_week', 'is_holiday', 'planned_tasks']]
y = data['completion_rate']

# 划分训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# 标准化特征
scaler = StandardScaler()
X_train_scaled = scaler.fit_transform(X_train)
X_test_scaled = scaler.transform(X_test)

# 构建TensorFlow模型
print("\n构建TensorFlow模型...")

# 定义模型
model = keras.Sequential([
    keras.layers.Dense(16, activation='relu', input_shape=(X_train_scaled.shape[1],)),
    keras.layers.Dense(8, activation='relu'),
    keras.layers.Dense(1)
])

# 编译模型
model.compile(optimizer='adam', loss='mse', metrics=['mae'])

# 显示模型摘要
model.summary()

# 训练模型
print("\n训练模型...")
history = model.fit(
    X_train_scaled, y_train,
    epochs=200,
    batch_size=4,
    validation_split=0.2,
    verbose=1,
    callbacks=[keras.callbacks.EarlyStopping(patience=20, restore_best_weights=True)]
)

# 绘制训练历史
plt.figure(figsize=(12, 5))
plt.subplot(1, 2, 1)
plt.plot(history.history['loss'])
plt.plot(history.history['val_loss'])
plt.title('模型损失')
plt.ylabel('损失')
plt.xlabel('轮次')
plt.legend(['训练', '验证'], loc='upper right')

plt.subplot(1, 2, 2)
plt.plot(history.history['mae'])
plt.plot(history.history['val_mae'])
plt.title('平均绝对误差')
plt.ylabel('MAE')
plt.xlabel('轮次')
plt.legend(['训练', '验证'], loc='upper right')
plt.savefig('output/training_history.png')

# 评估模型
print("\n评估模型...")
y_pred = model.predict(X_test_scaled)
mse = mean_squared_error(y_test, y_pred)
mae = mean_absolute_error(y_test, y_pred)
r2 = r2_score(y_test, y_pred)

print(f"均方误差 (MSE): {mse:.4f}")
print(f"平均绝对误差 (MAE): {mae:.4f}")
print(f"决定系数 (R²): {r2:.4f}")

# 可视化预测结果
plt.figure(figsize=(10, 6))
plt.scatter(y_test, y_pred)
plt.plot([y_test.min(), y_test.max()], [y_test.min(), y_test.max()], 'k--', lw=2)
plt.xlabel('实际完成率')
plt.ylabel('预测完成率')
plt.title('实际完成率 vs 预测完成率')
plt.savefig('output/actual_vs_predicted.png')

# 分析特征重要性
print("\n分析特征重要性...")

# 创建一个简单的模型来分析特征重要性
from sklearn.ensemble import RandomForestRegressor

rf_model = RandomForestRegressor(n_estimators=100, random_state=42)
rf_model.fit(X_train, y_train)

# 获取特征重要性
feature_importance = pd.DataFrame({
    'Feature': X.columns,
    'Importance': rf_model.feature_importances_
}).sort_values('Importance', ascending=False)

print("\n特征重要性:")
print(feature_importance)

# 可视化特征重要性
plt.figure(figsize=(10, 6))
sns.barplot(x='Importance', y='Feature', data=feature_importance)
plt.title('特征重要性')
plt.tight_layout()
plt.savefig('output/feature_importance.png')

# 分析不同特征值对完成率的影响
print("\n分析不同特征值对完成率的影响...")

# 1. 星期几对完成率的影响
day_of_week_impact = data.groupby('day_of_week')['completion_rate'].mean().reset_index()
print("\n星期几对完成率的影响:")
print(day_of_week_impact)

plt.figure(figsize=(10, 6))
sns.barplot(x='day_of_week', y='completion_rate', data=day_of_week_impact)
plt.title('星期几对完成率的平均影响')
plt.xlabel('星期几 (1=周一, 7=周日)')
plt.ylabel('平均完成率')
plt.savefig('output/day_of_week_impact.png')

# 2. 是否假日对完成率的影响
holiday_impact = data.groupby('is_holiday')['completion_rate'].mean().reset_index()
print("\n是否假日对完成率的影响:")
print(holiday_impact)

plt.figure(figsize=(10, 6))
sns.barplot(x='is_holiday', y='completion_rate', data=holiday_impact)
plt.title('是否假日对完成率的平均影响')
plt.xlabel('是否假日 (0=否, 1=是)')
plt.ylabel('平均完成率')
plt.savefig('output/holiday_impact.png')

# 3. 计划任务数量对完成率的影响
# 将计划任务数量分组
data['planned_tasks_group'] = pd.cut(data['planned_tasks'], bins=[4, 6, 8, 10, 12], labels=['5-6', '7-8', '9-10', '11-12'])
planned_tasks_impact = data.groupby('planned_tasks_group')['completion_rate'].mean().reset_index()
print("\n计划任务数量对完成率的影响:")
print(planned_tasks_impact)

plt.figure(figsize=(10, 6))
sns.barplot(x='planned_tasks_group', y='completion_rate', data=planned_tasks_impact)
plt.title('计划任务数量对完成率的平均影响')
plt.xlabel('计划任务数量')
plt.ylabel('平均完成率')
plt.savefig('output/planned_tasks_impact.png')

# 保存模型
model.save('output/completion_rate_model')
print("\n模型已保存到 'output/completion_rate_model'")

# 创建一个函数来预测新数据的完成率
def predict_completion_rate(day_of_week, is_holiday, planned_tasks):
    # 创建输入数据
    input_data = np.array([[day_of_week, is_holiday, planned_tasks]])
    
    # 标准化输入数据
    input_data_scaled = scaler.transform(input_data)
    
    # 预测完成率
    predicted_rate = model.predict(input_data_scaled)[0][0]
    
    return predicted_rate

# 示例：预测不同情况下的完成率
print("\n预测示例:")
scenarios = [
    {"day_of_week": 1, "is_holiday": 0, "planned_tasks": 8, "description": "周一，非假日，8个计划任务"},
    {"day_of_week": 5, "is_holiday": 0, "planned_tasks": 8, "description": "周五，非假日，8个计划任务"},
    {"day_of_week": 7, "is_holiday": 0, "planned_tasks": 8, "description": "周日，非假日，8个计划任务"},
    {"day_of_week": 3, "is_holiday": 1, "planned_tasks": 8, "description": "周三，假日，8个计划任务"},
    {"day_of_week": 3, "is_holiday": 0, "planned_tasks": 6, "description": "周三，非假日，6个计划任务"},
    {"day_of_week": 3, "is_holiday": 0, "planned_tasks": 10, "description": "周三，非假日，10个计划任务"},
    {"day_of_week": 3, "is_holiday": 0, "planned_tasks": 12, "description": "周三，非假日，12个计划任务"}
]

for scenario in scenarios:
    predicted_rate = predict_completion_rate(
        scenario["day_of_week"], 
        scenario["is_holiday"], 
        scenario["planned_tasks"]
    )
    print(f"{scenario['description']}: 预测完成率 = {predicted_rate:.4f}")

# 生成结论报告
print("\n生成结论报告...")

with open('output/conclusion_report.txt', 'w') as f:
    f.write("# 任务完成率影响因素分析报告\n\n")
    
    f.write("## 数据概述\n")
    f.write(f"- 总样本数: {len(data)}\n")
    f.write(f"- 平均完成率: {data['completion_rate'].mean():.4f}\n")
    f.write(f"- 完成率标准差: {data['completion_rate'].std():.4f}\n")
    f.write(f"- 最高完成率: {data['completion_rate'].max():.4f}\n")
    f.write(f"- 最低完成率: {data['completion_rate'].min():.4f}\n\n")
    
    f.write("## 模型性能\n")
    f.write(f"- 均方误差 (MSE): {mse:.4f}\n")
    f.write(f"- 平均绝对误差 (MAE): {mae:.4f}\n")
    f.write(f"- 决定系数 (R²): {r2:.4f}\n\n")
    
    f.write("## 特征重要性\n")
    for _, row in feature_importance.iterrows():
        f.write(f"- {row['Feature']}: {row['Importance']:.4f}\n")
    f.write("\n")
    
    f.write("## 星期几对完成率的影响\n")
    for _, row in day_of_week_impact.iterrows():
        f.write(f"- 星期{int(row['day_of_week'])}: {row['completion_rate']:.4f}\n")
    f.write("\n")
    
    f.write("## 是否假日对完成率的影响\n")
    for _, row in holiday_impact.iterrows():
        holiday_status = "是" if row['is_holiday'] == 1 else "否"
        f.write(f"- 假日({holiday_status}): {row['completion_rate']:.4f}\n")
    f.write("\n")
    
    f.write("## 计划任务数量对完成率的影响\n")
    for _, row in planned_tasks_impact.iterrows():
        f.write(f"- {row['planned_tasks_group']}个任务: {row['completion_rate']:.4f}\n")
    f.write("\n")
    
    f.write("## 主要发现